#include "Polynomicon.h"
#include "FourierTron.h"
#include "VectorOps.h"
#include <iostream>


using namespace std;
using namespace SeamStress;


void Polynomicon::multiplyPoly(const vector<cxfloat>& in1, const vector<cxfloat>& in2, vector<cxfloat>& out)
{
  tempvec1.clear();
  tempvec1 = in1;
  unsigned long int MM = in1.size() + in2.size() - 1;
  unsigned long int N = 1;
  while(MM>=N){N<<=1;}
  tempval1.set(0,0,prec);
  tempvec1.resize(N, tempval1);
  VectorOps vecops(pins.seamstresses);
  if( out.size() < N )
  {
    out.resize(N);
    vecops.setVector(&out, tempval1);
  }
  
  vecops.copyVector(&in2, &out);
  
  fourier->CT_FFT(tempvec1, N);
  fourier->CT_FFT(out, N);
  
  thread_var1=N;
  if(N<nthreads){thread_tot=N;}
  else{thread_tot=nthreads;}
  tempvecpointer1 = &out;
  pins.sewStraight(&Polynomicon::multiplyPolyThread1, thread_tot);
  
  fourier->inverseCT_FFT(out, N);
  if(out.size() > MM){out.resize(MM, tempval1);}
}


void Polynomicon::multiplyPolyThread1(void* arg)
{
  unsigned long int w = (*((unsigned long int *)arg));
  
  unsigned long int i = w;
  
  unsigned long int max = thread_var1 - thread_tot;
  while(i < max)
  {
    (*tempvecpointer1)[i].mul(tempvec1[i]);
    
    i += thread_tot;
  }
  (*tempvecpointer1)[i].mul(tempvec1[i]);
}


void Polynomicon::multiplyPoly(const refloat_vec& in1, const refloat_vec& in2, refloat_vec& out)
{
  unsigned long int in1size = in1.size();
  unsigned long int in2size = in2.size();
  
  if( (in1size <= 16) && (in2size <= 16) )
  {
    retempval1.set(0);
    out.clear();out.assign( in1size+in2size-1, retempval1 );
    retempvecpointer1 = &out;
    retempvecpointer2 = &in1;
    retempvecpointer3 = &in2;
    
    pins.sewStraight(&Polynomicon::multiplyPolySchoolbook, nthreads);
    return;
  }
  
  unsigned long int MM = in1.size() + in2.size() - 1;
  unsigned long int N = 1;
  while(MM>=N){N<<=1;}
  
  VectorOps vecops(pins.seamstresses);
  refloat_vec vec_temp_1;vec_temp_1.assign(N, nullptr);
  retempval1.set(0);vecops.setVector(&vec_temp_1, retempval1);
  vecops.copyVector(&in1, &vec_temp_1);
  refloat_vec vec_temp_2;vec_temp_2.assign(N, nullptr);
  retempval1.set(0);vecops.setVector(&vec_temp_2, retempval1);
  vecops.copyVector(&in2, &vec_temp_2);
  
  fourier->CT_FFT(vec_temp_1, tempvec1, N);
  fourier->CT_FFT(vec_temp_2, tempvec2, N);
  
  thread_var1=N;
  if(N<nthreads){thread_tot=N;}
  else{thread_tot=nthreads;}
  tempvecpointer1 = &tempvec2;
  pins.sewStraight(&Polynomicon::multiplyPolyThread1, thread_tot);
  
  fourier->inverseCT_FFT(tempvec2, out, N);
  if(out.size() > MM){out.resize(MM, retempval1);}
}


void Polynomicon::multiplyPolySchoolbook(void* arg)
{
  unsigned long int w = (*((unsigned long int *)arg));
  for(unsigned long int i=w,loopsizei=retempvecpointer1->size();i<loopsizei;i+=nthreads)
  {
    for(unsigned long int j=0,loopsizej=i;j<=loopsizej;j+=1)
    {
      if( (j >= retempvecpointer2->size()) || ((i-j) >= retempvecpointer3->size()) ){continue;}
      (*retempvecpointer1)[i] += (*retempvecpointer2)[j] * (*retempvecpointer3)[i-j];
    }
  }
}


void Polynomicon::multiplyPoly(const refloat_vec& in1, const refloat_vec& in2, refloat_vec& out, long int log2scale)
{
  VectorOps vecops(pins.seamstresses);
  refloat_vec vec_temp_1;vec_temp_1.assign(in1.size(), nullptr);
  vecops.copyVector(&in1, &vec_temp_1);
  refloat_vec vec_temp_2;vec_temp_2.assign(in2.size(), nullptr);
  vecops.copyVector(&in2, &vec_temp_2);
  
  p2_scale = log2scale;
  
  thread_var1 = vec_temp_1.size();
  if(thread_var1<nthreads){thread_tot=thread_var1;}
  else{thread_tot=nthreads;}
  retempvecpointer1 = &vec_temp_1;
  pins.sewStraight(&Polynomicon::scalePolyPower2Real, thread_tot);
  
  thread_var1 = vec_temp_2.size();
  if(thread_var1<nthreads){thread_tot=thread_var1;}
  else{thread_tot=nthreads;}
  retempvecpointer1 = &vec_temp_2;
  pins.sewStraight(&Polynomicon::scalePolyPower2Real, thread_tot);
  
  multiplyPoly(vec_temp_1, vec_temp_2, out);
  
  p2_scale = -log2scale;
  
  thread_var1 = vec_temp_1.size();
  if(thread_var1<nthreads){thread_tot=thread_var1;}
  else{thread_tot=nthreads;}
  retempvecpointer1 = &vec_temp_1;
  pins.sewStraight(&Polynomicon::scalePolyPower2Real, thread_tot);
  
  thread_var1 = vec_temp_2.size();
  if(thread_var1<nthreads){thread_tot=thread_var1;}
  else{thread_tot=nthreads;}
  retempvecpointer1 = &vec_temp_2;
  pins.sewStraight(&Polynomicon::scalePolyPower2Real, thread_tot);
  
  thread_var1 = out.size();
  if(thread_var1<nthreads){thread_tot=thread_var1;}
  else{thread_tot=nthreads;}
  retempvecpointer1 = &out;
  pins.sewStraight(&Polynomicon::scalePolyPower2Real, thread_tot);
}


void Polynomicon::scalePoly(vector<cxfloat>& poly, const cxfloat& scale)
{
  if(poly.size() < 4)
  {
    tempval1.set(1,0);
    for(unsigned long int i=0;i<poly.size();i++)
    {
      poly[i].mul(tempval1);
      tempval1.mul(scale);
    }
    return;
  }
  
  if(nthreads <= poly.size()){thread_tot = nthreads;}
  else{thread_tot = poly.size();}
  
  tempval2.set(scale);
  pow(thread_tot, tempval2, tempval1);
  tempvecpointer1 = &poly;
  thread_var1 = poly.size();
  pins.sewStraight(&Polynomicon::powerUp, thread_tot);
}


void Polynomicon::powerUp(void* arg)
{
  unsigned long int w = (*((unsigned long int *)arg));
  unsigned long int max = thread_var1 - thread_tot;
  unsigned long int i = w;
//   __builtin_prefetch(&((*tempvecpointer1)[i]), 1, 3);
  pow(w, tempval2, thread_vec1[w]);
  while(i < max)
  {
//     __builtin_prefetch(&((*tempvecpointer1)[i+thread_tot]), 1, 3);
    (*tempvecpointer1)[i].mul(thread_vec1[w]);
    thread_vec1[w].mul(tempval1);
    i += thread_tot;
  }
  (*tempvecpointer1)[i].mul(thread_vec1[w]);
  thread_vec1[w].mul(tempval1);
}


void Polynomicon::scalePoly(refloat_vec& poly, const refloat& scale)
{
  if(poly.size() < 4)
  {
    retempval1.set(1);
    for(unsigned long int i=0;i<poly.size();i++)
    {
      poly[i].mul(retempval1);
      retempval1.mul(scale);
    }
    return;
  }
  
  if(nthreads <= poly.size()){thread_tot = nthreads;}
  else{thread_tot = poly.size();}
  
  retempval2.set(scale);
  pow(thread_tot, retempval2, retempval1);
  retempvecpointer1 = &poly;
  thread_var1 = poly.size();
  pins.sewStraight(&Polynomicon::powerUpReal, thread_tot);
}


void Polynomicon::powerUpReal(void* arg)
{
  unsigned long int w = (*((unsigned long int *)arg));
  unsigned long int max = thread_var1 - thread_tot;
  unsigned long int i = w;
//   __builtin_prefetch(&((*retempvecpointer1)[i]), 1, 3);
  pow(w, retempval2, rethread_vec1[w]);
  while(i < max)
  {
//     __builtin_prefetch(&((*retempvecpointer1)[i+thread_tot]), 1, 3);
    (*retempvecpointer1)[i].mul(rethread_vec1[w]);
    rethread_vec1[w].mul(retempval1);
    i += thread_tot;
  }
  (*retempvecpointer1)[i].mul(rethread_vec1[w]);
  rethread_vec1[w].mul(retempval1);
}


void Polynomicon::evaluatePoly(refloat& output, const refloat_vec& poly, const refloat& input)
{
  if(poly.size() == 0){output.set(0);return;}
  else if(poly.size() < 4)
  {
    output.set(poly.back());
    for(unsigned long int i=(poly.size()-1);i>0;i--)
    {
      output.mul(input);
      output.add(poly[i-1]);
    }
    return;
  }
  
  if(nthreads <= poly.size()){thread_tot = nthreads;}
  else{thread_tot = poly.size();}
  
  retempval2.set(input);
  pow(thread_tot, retempval2, retempval1);
  retempvecpointer2 = &poly;
  thread_var1 = poly.size();
  pins.sewStraight(&Polynomicon::parallelHornerReal, thread_tot);

  unsigned long int pow2threads = 1;
  while(pow2threads < thread_tot)
  {
    pow2threads<<=1;
  }
  if(pow2threads > thread_tot){pow2threads>>=1;}
  output.set(0);
  for(unsigned long int i=pow2threads;i<thread_tot;i++)
  {
    output.add(rethread_vec1[i]);
  }
  thread_tot = (pow2threads>>1);
  while(thread_tot > 0)
  {
    pins.sewStraight(&Polynomicon::addIntermediate, thread_tot);
    thread_tot>>=1;
  }
  output.add(rethread_vec1[0]);
}


void Polynomicon::addIntermediate(void* arg)
{
  unsigned long int w = (*((unsigned long int *)arg));
  rethread_vec1[w].add(rethread_vec1[w+thread_tot]);
}


void Polynomicon::parallelHornerReal(void* arg)
{
  unsigned long int w = (*((unsigned long int *)arg));
  unsigned long int mod = (retempvecpointer2->size()-1)%(thread_tot);
  unsigned long int base = (retempvecpointer2->size()-1) - mod;
  unsigned long int backstart = 0;
  if(w<=mod)
  {
    backstart = base + w;
  }
  else
  {
    backstart = base + w - thread_tot;
  }
  rethread_vec1[w].set((*retempvecpointer2)[backstart]);
  for(unsigned long int i=backstart;i>=thread_tot;i-=thread_tot)
  {
    rethread_vec1[w].mul(retempval1);
    rethread_vec1[w].add((*retempvecpointer2)[i-thread_tot]);
  }
  pow(w, retempval2, rethread_vec2[w]);
  rethread_vec1[w].mul(rethread_vec2[w]);
}


void Polynomicon::vectorMultiply(vector<cxfloat>& poly, const cxfloat& scale)
{
  tempval1.set(scale);
  tempvecpointer1 = &poly;
  if(poly.size() >= nthreads){thread_tot=nthreads;}
  else{thread_tot = poly.size();}
  pins.sewStraight(&Polynomicon::vectorMultiply, thread_tot);
}


void Polynomicon::vectorMultiply(void* arg)
{
  unsigned long int w = (*((unsigned long int *)arg));
  
  unsigned long int i = w;
  
  unsigned long int max = (*tempvecpointer1).size() - thread_tot;
  
//   __builtin_prefetch(&((*tempvecpointer1)[i]), 1, 3);
  while(i < max)
  {
//     __builtin_prefetch(&((*tempvecpointer1)[i+thread_tot]), 1, 3);
    (*tempvecpointer1)[i].mul(tempval1);
    
    i += thread_tot;
  }
  (*tempvecpointer1)[i].mul(tempval1);
}


void Polynomicon::vectorMultiply(refloat_vec& poly, const refloat& scale)
{
  retempval1.set(scale);
  retempvecpointer1 = &poly;
  if(poly.size() >= nthreads){thread_tot=nthreads;}
  else{thread_tot = poly.size();}
  pins.sewStraight(&Polynomicon::vectorMultiplyReal, thread_tot);
}


void Polynomicon::vectorMultiplyReal(void* arg)
{
  unsigned long int w = (*((unsigned long int *)arg));
  
  unsigned long int i = w;
  
  unsigned long int max = (*retempvecpointer1).size() - thread_tot;
  
//   __builtin_prefetch(&((*retempvecpointer1)[i]), 1, 3);
  while(i < max)
  {
//     __builtin_prefetch(&((*retempvecpointer1)[i+thread_tot]), 1, 3);
    (*retempvecpointer1)[i].mul(retempval1);
    
    i += thread_tot;
  }
  (*retempvecpointer1)[i].mul(retempval1);
}


//output[j] = Sum_{k=j}^{n} vec1[k] * vec2[k-j] * vec3[j]
//done in O(n*log(n)) operations
//log2scale allows a temporary change of variables which can improve numerical stability
void Polynomicon::generalizedShift(refloat_vec& vec1, refloat_vec& vec2, refloat_vec& vec3, refloat_vec& output, long int log2scale)
{
  if(nthreads > ((vec1.size() - 1) >> 1)){thread_tot = ((vec1.size() - 1) >> 1);}
  else{thread_tot = nthreads;}
  thread_var1 = ((vec1.size() - 1) >> 1);
  thread_var2 = (vec1.size() - 1);
  retempvecpointer1 = &vec1;
  pins.sewStraight(&Polynomicon::reflectAboutCenter, thread_tot);
  if(thread_var1 != (thread_var2 - thread_var1))
  {
    retempval1.set(vec1[thread_var1]);
    vec1[thread_var1].set(vec1[thread_var2 - thread_var1]);
    vec1[thread_var2 - thread_var1].set(retempval1);
  }
  
  multiplyPoly(vec1, vec2, output, log2scale);
  output.resize(vec1.size(), retempval1);
  
  if(nthreads > ((vec1.size() - 1) >> 1)){thread_tot = ((vec1.size() - 1) >> 1);}
  else{thread_tot = nthreads;}
  thread_var1 = ((vec1.size() - 1) >> 1);
  thread_var2 = (vec1.size() - 1);
  retempvecpointer1 = &output;
  pins.sewStraight(&Polynomicon::reflectAboutCenter, thread_tot);
  if(thread_var1 != (thread_var2 - thread_var1))
  {
    retempval1.set(output[thread_var1]);
    output[thread_var1].set(output[thread_var2 - thread_var1]);
    output[thread_var2 - thread_var1].set(retempval1);
  }
  
  retempvecpointer1 = &output;
  retempvecpointer2 = &vec3;
  if(nthreads > output.size()){thread_tot = output.size();}
  else{thread_tot = nthreads;}
  thread_var1 = output.size();
  pins.sewStraight(&Polynomicon::oneToOneMultiply, thread_tot);
}


void Polynomicon::reflectAboutCenter(void* arg)
{
  unsigned long int w = (*((unsigned long int *)arg));
  
  for(unsigned long int i=w;i<thread_var1;i+=thread_tot)
  {
    rethread_vec1[w].set((*retempvecpointer1)[i]);
    (*retempvecpointer1)[i].set((*retempvecpointer1)[thread_var2 - i]);
    (*retempvecpointer1)[thread_var2 - i].set(rethread_vec1[w]);
  }
}


void Polynomicon::oneToOneMultiply(void* arg)
{
  unsigned long int w = (*((unsigned long int *)arg));
  
  for(unsigned long int i=w;i<thread_var1;i+=thread_tot)
  {
    (*retempvecpointer1)[i].mul((*retempvecpointer2)[i]);
  }
}

